# -*- coding: utf-8 -*-
"""
Created on Tue Apr 24 13:33:37 2018

@author: maowenxiang
"""

import re

def check_sentence(sentence):
    def contain_digit(sentence):
        pattern = re.compile('[0-9]+')
        return True if re.findall(pattern, sentence) else False
    return not sentence.islower() or contain_digit(sentence)

def extract_question(line):
    return line.split(">	")[1].strip().replace("’s", "").replace("'s","")

def get_answer_entities(answer):
	entity_list = []
	for index, part in enumerate(answer.split('mso:')):
		if index > 0:
			entity_candidates = part.split(' ')
			for entity in filter(lambda x: not x.startswith('?'), [entity_candidates[1], entity_candidates[2]]):
				entity_list.append(entity.replace('_', ' '))
                
	return entity_list

def check_is_num(entity):
    return set(entity).issubset(set(list(map(str, range(10))) + ['.']))

def check_is_date(entity):
    return set(entity).issubset(set(list(map(str, range(10))) + ['/']))

def update_question(question, entity_list):
    for index, entity in enumerate(entity_list, 1):
        replace_string = 'ENTITY' + str(index)
        if check_is_num(entity):
            replace_string = 'NUM'
        elif check_is_date(entity):
            replace_string = 'DATE'
        question = question.replace(entity, replace_string)
    return question
def predict_question_entities(question):
    entity_list, tokens = [], question.split()
    index, length = 0, len(tokens)
    word_set = {'get', 'got', 'had', 'have', 'won', 'is'}
    def check_token(token):
        return token[0].isupper() or token[0].isdigit() or token.startswith(('(', '{', '.', '-')) or \
               (len(token) > 1 and token[1].isupper()) or (not token.isalpha() and '-' not in token) or \
                any(ord(char) > 240 for char in token) 
    def next_short_token_check(index):
        token = tokens[index]
        if token in word_set or (token in ['from', 'in', 'since'] and index + 1 < length and tokens[index+1][0].isdigit()):
            return False
        return (len(token) <= 4 and (index + 1 < length) and check_token(tokens[index+1])) or \
               (token in ['and', 'to', 'in', '&', 'of', 'at', 'with', 'for', 'on', 'from', 'de', 'pour'] and \
                (index + 1 < length) and tokens[index + 1] in ['the', 'a', 'to', 'la'] and (index + 2 < length) \
                and check_token(tokens[index + 2]))       
    while index < length:
        entity = []
        while index < length and (check_token(tokens[index]) or next_short_token_check(index)):
            if len(entity) == 0 and check_token(tokens[index]) == False:
                index += 1	            
                continue       			          
            entity.append(tokens[index])
            index += 1
        if not entity:
            index += 1
        else:
            entity_list.append(' '.join(entity))
    return entity_list

def predict_and_evaluate(train_file, output_file):
    total, error = 0, 0  
    f_out = open(output_file, 'w', encoding='utf-8')
    with open(train_file, encoding="utf-8") as f:
        lines = f.readlines()
    for i, line in enumerate(lines):
        if i % 3 == 0:
            question = extract_question(line)
            answer = lines[i+1].split('\t')[1].replace("’s", "").replace("'s", "")
            total += 1
            predicts = predict_question_entities(question)
            labels = get_answer_entities(answer)
            if set(predicts) != set(labels):
                print('问题:{} \n 预测: {} 实际:{}'.format(question, predicts, labels))
                error += 1
            if check_sentence(question):
                entity_list = predict_question_entities(question)
            else:
                entity_list = ['$$$全由小写字母组成，无法预测实体']
            f_out.write(str(i//3 + 1) + '---' + update_question(question, entity_list) + '\n')            
    f_out.close()            
    print('错误个数:{} 总数:{} 正确率:{:.2f}%'.format(error, total, 100 * (total - error) / total))

def predict(test_file, output_file, cannot_predict_file):
    cannot_predict, total = 0, 0
    with open(test_file, encoding='utf-8') as f_in, open(output_file, 'w', encoding='utf-8') as f_out,\
                        open(cannot_predict_file, 'w', encoding='utf-8') as f_cannot_predict:
        for line in f_in:
            if line[0] == '<':
                total += 1
                question = extract_question(line)
                if check_sentence(question):
                    entity_list = predict_question_entities(question)
                else:
                    entity_list = ['$$$全由小写字母组成，无法预测实体']
                    f_cannot_predict.write(question + '\n')
                    cannot_predict += 1
                f_out.write(question + '\n' + '***'.join(entity_list) + '\n' + update_question(question, entity_list))
    print('问题总数: {} 无法预测问题的实体的问题个: {}'.format(total, cannot_predict))
    
predict_and_evaluate('train.txt', 'predict_traindata_entities.txt')

predict('test.txt', 'predict_testdata_entities.txt', 'cannot_predict.txt')